from PIL import Image
import numpy as np
import tensorflow as tf

model_save_path = './18 checkpoint/mnist.ckpt'

model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(128, activation='relu'),
    tf.keras.layers.Dense(10, activation='softmax')
])

model.load_weights(model_save_path)

while(1):
    image = input('the test number:')
    image_path = './20 mnist/'+ image + '.png'
    img = Image.open(image_path)
    # 预处理 
    img = img.resize((28,28),Image.ANTIALIAS)
    img_arr = np.array(img.convert('L'))

    img_arr = 255 - img_arr

    img_arr = img_arr/255.0
    x_predict = img_arr[tf.newaxis, ... ]
    result = model.predict(x_predict)
    pred = tf.argmax(result,axis=1)
    tf.print(pred)
    print('\n')